# -*- coding: utf-8 -*-
import os
import math
from jsonpath import jsonpath
from locust import LoadTestShape
from locust import HttpUser, task, constant


class QuickStartUser(HttpUser):
    host = "http://127.0.0.1:8000"
    wait_time = constant(1)

    @task
    def api(self):
        with self.client.get(url="/locust_test", catch_response=True, name='测试') as response:
            if response.status_code == 200:
                try:
                    assert jsonpath(response.json(), '$..code')[0] == 1
                except:
                    response.failure("Timeout")
            else:
                response.failure(response.status_code)


class StagesShape(LoadTestShape):
    """
        step_time -- 逐步加载时间长度
        step_load -- 用户每一步增加的量
        spawn_rate -- 用户在每一步的停止/启动的多少用户数
        time_limit -- 时间限制压测的执行时长
    """

    # 逐步负载策略每隔30秒新增启动10个用户
    setp_time = 30
    setp_load = 10
    spawn_rate = 10
    time_limit = 300

    def tick(self):
        run_time = self.get_run_time()

        if run_time > self.time_limit:
            return None
        current_step = math.floor(run_time / self.setp_time) + 1
        print(current_step)
        tick_data = (current_step * self.setp_load, self.spawn_rate)
        return tick_data


if __name__ == "__main__":
    os.system("locust -f 负载图_逐步负载图策略.py --web-host=127.0.0.1")
